#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@gmail.com
@time: 2018/3/25 10:20
@desc: 
"""

import os
import time

import cv2
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from dltools.train.tools import load_model, merge_bounding_boxes, run_detection
from dltools.utils.basic import get_name_and_ext
from dltools.utils.log import get_console_logger
from object_detection.utils import visualization_utils as vis_util


def predict_image(checkpoint, data_dir):
    start_time = time.time()
    model = load_model(checkpoint)
    stage_time = time.time()
    logger.info('{} elapsed time: {:.3f}s'.format(time.ctime(),
                                                  stage_time - start_time))
    config = tf.ConfigProto(device_count={"CPU": 4, "GPU": 1})
    with model.as_default():
        with tf.Session(graph=model, config=config) as sess:
            with tf.device("/gpu:0"):
                for idx, image_name in enumerate(os.listdir(data_dir)):
                    if image_name.split('.')[-1] == 'png' or \
                            image_name.split('.')[-1] == 'jpg':
                        logger.info('predicting image: {}!'.format(
                            os.path.join(data_dir, image_name)))
                        img_head_name, _ = get_name_and_ext(data_dir)
                        img = cv2.imread(os.path.join(data_dir, image_name))
                        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                        # img = cv2.resize(img, (500, 500))
                        # img = (2.0 - 255.0) * img - 1.0
                        input_image = np.expand_dims(img, 0)
                        _boxes, _classes, _scores = run_detection(sess, model,
                                                                  input_image)
                        stage_time = time.time()
                        logger.info(
                            'elapsed time: {:.3f}s'.format(stage_time - start_time))
                        logger.info('predicting image completed !')
                        start_time = stage_time
                        boxes, classes, scores = [], [], []
                        for box, cls, sc in zip(_boxes, _classes, _scores):
                            if sc >= 0.5:
                                boxes.append(box)
                                classes.append(cls)
                                scores.append(sc)
                        print(scores)
                        # _boxes, _classes, _scores = merge_bounding_boxes(
                        #     np.array(_boxes), np.array(_scores),
                        #     np.array(_classes),
                        #     0.8)
                        _boxes = np.array(_boxes)
                        _classes = np.array(_classes).astype(np.int32)
                        _scores = np.array(_scores)
                        vis_util.visualize_boxes_and_labels_on_image_array(
                            img, _boxes, _classes, _scores,
                            {1: {'id': 1, 'name': 'p'}},
                            use_normalized_coordinates=True,
                            min_score_thresh=0.5,
                            line_thickness=5)
                        plt.imshow(img)
                        plt.show()


if __name__ == '__main__':
    logger = get_console_logger('Predict Task')
    predict_image('D:\\workspace\\data\\car\\trained\\export\\frozen_inference_graph.pb',
                  'D:\\workspace\\data\\car\\test_image\\1131')
                  # 'D:\\workspace\\data\\car\\test_image')
